'''
Author: caishuyang
Date: 2023-03-13 21:09:52
LastEditors: caishuyang
LastEditTime: 2023-03-13 21:39:04
Description: 训练模型
'''
import torch.utils.data as Data  
from _02PipeDatasetLoader import PipeDataset
from torchvision.transforms import transforms
import torch
from torch import nn
from _03Deeplabv3plus import DeepLabV3
import os

trans=transforms.ToTensor()
datapath="Dataset\\Train"
pipe=PipeDataset(datapath,trans,trans)
trainiter=Data.DataLoader(pipe,4,shuffle=True,num_workers=0,drop_last=True)
modelpath=os.path.join("Model","model_01.pth") #设定模型存储路径
torch.cuda.set_device(0)
net=DeepLabV3(class_num=2) #实现车道线和背景的二分类
net=net.to("cuda") #选用cuda加速
Criterion = nn.BCELoss().to('cuda')
Optimizer = torch.optim.Adam(net.parameters(), lr=0.01) #学习率设为0.01
epochs=1000 #训练轮数

'''
训练过程，为了省事，没有加验证环节
'''

import gc
for epoch in range(epochs):
    net.train()     
    gc.collect()
    torch.cuda.empty_cache()
    loss=0
    for data in trainiter:
        
        img,label=data
        imgg=img.to("cuda")
        labelg=label.to("cuda")
        Optimizer.zero_grad()
        with torch.set_grad_enabled(True):
            OutputImg = net(imgg)
            #print(OutputImg.size())
            #print(label.size())
            BatchLoss = Criterion(OutputImg, labelg)
            BatchLoss.backward(retain_graph=False)
            Optimizer.step()
            loss+=BatchLoss.item()
    print(loss)

torch.save(net.state_dict(),modelpath)